# encoding: utf-8
# !/usr/bin/python3
import torch.nn as nn
import math
import torch
import torch.nn.functional as F





def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class double_conv(nn.Module):
    '''(conv => BN => ReLU) * 2'''
    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            # nn.BatchNorm2d(out_ch),
            nn.GroupNorm(4,out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            # nn.BatchNorm2d(out_ch),
            nn.GroupNorm(4,out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class inconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(inconv, self).__init__()
        self.conv = double_conv(in_ch, out_ch)

    def forward(self, x):
        x = self.conv(x)
        return x



class up(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(up, self).__init__()


        if bilinear:

            self.up = nn.Sequential(nn.Conv2d(in_ch,in_ch//2,1),
                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True))

        else:
            self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)

        self.conv = double_conv(in_ch, out_ch)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffX = x1.size()[2] - x2.size()[2]
        diffY = x1.size()[3] - x2.size()[3]
        x2 = F.pad(x2, (diffX // 2, int(diffX / 2),
                        diffY // 2, int(diffY / 2)))
        # print('x1',x1.shape,'x2',x2.shape)
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


class Aspp(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(Aspp, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=3,dilation=1,padding=1)
        self.conv2 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=3,dilation=2,padding=2)
        self.conv4 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=3,dilation=4,padding=4)
        self.conv8 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=3,dilation=8,padding=8)
        self.conv_all = nn.Conv2d(4*out_channel,out_channel,kernel_size=1)
    def forward(self,input):
        conv1 = self.conv1(input)
        conv2 = self.conv2(input)
        conv4 = self.conv4(input)
        conv8 = self.conv8(input)
        all_conv = torch.cat((conv1,conv2,conv4,conv8),dim = 1)
        all_conv = self.conv_all(all_conv)
        return all_conv

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.GroupNorm(4,planes)
        # self.bn1 = nn.InstanceNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.GroupNorm(4,planes)
        # self.bn2 = nn.InstanceNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class outconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(outconv, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 1)

    def forward(self, x):
        x = self.conv(x)
        return x



class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes,planes, kernel_size=1, bias=False)
        self.bn1 = nn.GroupNorm(4,planes)
        # self.bn1 = nn.InstanceNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.GroupNorm(4,planes)
        # self.bn2 = nn.InstanceNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.GroupNorm(4,planes * self.expansion)
        # self.bn3 = nn.InstanceNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ResNet_Unet(nn.Module):

    def __init__(self, block = BasicBlock, layers = [2,2,2,2], num_classes=7,use_aspp = False):
        self.inplanes = 64
        self.use_aspp = use_aspp
        super(ResNet_Unet, self).__init__()
        self.conv_first = nn.Conv2d(4, 64, kernel_size=7, stride=1, padding=3,
                               bias=False)
        self.bn1 = nn.GroupNorm(4,64)
        # self.bn1 = nn.InstanceNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.conv1_add = nn.Conv2d(64,32,1)
        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0],stride=2)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.conv_middle = double_conv(512,1024)
        if self.use_aspp:
            self.aspp_layer = Aspp(in_channel=512,out_channel=512)
        self.up1 = up(512, 256)
        self.up2 = up(256, 128)
        self.up3 = up(128, 64)
        self.up4 = up(64, 32)
        self.outc = outconv(32, num_classes)

        #参数初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            # elif isinstance(m, nn.InstanceNorm2d):
            #     nn.init.constant_(m.weight, 1)
            #     nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                # nn.BatchNorm2d(planes * block.expansion),
                nn.GroupNorm(4,planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x0 = self.conv_first(x)
        # print('conv1',x0.shape)
        # conv1[64,256,256]
        x0 = self.bn1(x0)
        x0 = self.relu(x0)

        # x = self.maxpool(x)


        x1 = self.layer1(x0)
        # print('layer1',x1.shape)

        x2 = self.layer2(x1)
        # print('layer2',x2.shape)
        x3 = self.layer3(x2)
        # print('layer3',x3.shape)
        x4 = self.layer4(x3)
        if self.use_aspp:
            x4 = self.aspp_layer(x4)
        # print('layer4',x4.shape)
        x = self.up1(x4, x3)
        x = self.up2(x, x2)
        x = self.up3(x, x1)
        # print('xxx',x.shape)
        x0 = self.conv1_add(x0)
        x = self.up4(x, x0)
        x = self.outc(x)


        return x


def load_pre_model_dict(model,pre_train_dir):
    pretrain_dict = torch.load(pre_train_dir)

    model_dict = model.state_dict()

    pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and 'last_linear' not in k}
    model_dict.update(pretrain_dict)
    # print(pretrain_dict.keys())
    model.load_state_dict(model_dict)
    # print('load pre_trained_models')


def unet_resnet_model_GN(layers,num_classes,use_aspp = False,pre_train = False,):
    backbone = {
        'resnet18':[2, 2, 2, 2],
        'resnet34':[3, 4, 6, 3],
        'resnet50':[3, 4, 6, 3],
        'resnet101':[3, 4, 23, 3],
        'resnet152':[3, 8, 36, 3]
    }
    pretrain_file = {
        'resnet18':'../pre_trained_models/resnet18.pth',
        'resnet34':'../pre_trained_models/resnet34.pth',
        'resnet50':'../pre_trained_models/resnet50.pth',
        'resnet101':'../pre_trained_models/resnet101.pth'
    }
    model = ResNet_Unet(layers=backbone[layers],num_classes=num_classes,use_aspp=use_aspp).cuda()
    # 载入resnet18的预训练模型。
    if pre_train:
        load_pre_model_dict(model, pre_train_dir=pretrain_file[layers])
    return model
if __name__ == '__main__':
    from torchsummary import summary
    model = unet_resnet_model_GN(layers='resnet101',num_classes=2,use_aspp=True,pre_train=False)
    summary(model,(3,320,512))
